#!/usr/bin/python
# -*- coding: utf-8 -*-

# Copyright (c) 2011
#
# Permission is hereby granted, free of charge, to any person obtaining a
# copy of this software and associated documentation files (the "Software"),
# to deal in the Software without restriction, including without limitation

# the rights to use, copy, modify, merge, publish, distribute, sublicense,
# and/or sell copies of the Software, and to permit persons to whom the
# Software is furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.
#
# Author: Jesus Carrero <j.o.carrero@gmail.com>
# Mountain View, CA
#

from itertools import izip, count

from numpy import meshgrid

from scipy.linalg import norm
from scipy import array, exp, empty, asarray
from scipy import linspace, ones, r_, c_

from models.CalibVolCurve import CalibVolCurve
from models.CalibVolSurfCCE import CalibVolSurfCCE
from models.CalibrateD_VolSurface import CalibrateDVS

from models.AnalyticalFormulas import vanillaCallPrice, \
    vanillaPutPrice, vegaVanilla

from input_output.DerivativePlots import DerivativePlots

from utils.MainLog import runGlobalLogger
import unittest


class CalVolTestSuite(unittest.TestCase):

    def setUp(self):

        # market data

        self.m_r = float(0.05)
        self.q = 0.0

        # contract description :)

        sprices = asarray([[40], [35], [40], [40]])
        self.m_strikes = asarray([30, 30, 45, 30])
        self.m_prices = asarray([float(S) for S in sprices])
        self.m_cstyles = ['E', 'E', 'E', 'E']
        self.m_contract = ['C', 'C', 'C', 'C']
        self.m_t2mat = asarray([.4, .2, 1., .2])

        # market value

        self.sigmas = asarray([.4, .5, .3, .2])
        call_prices = asarray([vanillaCallPrice(S0, K, self.m_r, sg,
                              tm) for (S0, K, sg, tm) in zip(sprices,
                              self.m_strikes, self.sigmas,
                              self.m_t2mat)])

        put_prices = asarray([vanillaPutPrice(S0, K, self.m_r, sg, tm)
                             for (S0, K, sg, tm) in zip(sprices,
                             self.m_strikes, self.sigmas,
                             self.m_t2mat)])

        self.exact_sol = r_[put_prices[0:2].flatten(),
                           call_prices[2:].flatten()]
        self.exact_sol = call_prices

        # optimizatio engine

        self.m_regularization = .5
        self.max_num_iter = 50
        self.m_random_trails = False

        # numerical PDE

        self.ntime_steps = 200
        self.m_nsegs = 20


class CalCurve(CalVolTestSuite):
    """ calibrate volatility curve. """

    def runTest(self):
        """ 1. - instantiate the engine.
            2. - set calibration data.
            3. - set initial guess.
            4. - set calibration method.
            5. - calibrate.
        """

        calb_vol_curve = CalibVolCurve()
        calb_vol_curve.set_grid_spacing(self.ntime_steps, self.m_nsegs)

        (ki, si, r) = (self.m_strikes, self.m_prices, self.m_r)
        (ti, style) = (self.m_t2mat, self.m_cstyles)
        calb_vol_curve.set_cal_data(ki, si, r, ti, style, self.m_contract)
        calb_vol_curve.set_market_value(self.exact_sol)

        # constant inital volatility.

        calb_vol_curve.set_initial_guess([.1])

        calb_vol_curve.set_cal_method('fmin')
        (m_t2M, sig) = calb_vol_curve.calibrate()
        print 'initial Guess ', pow(.2, 2), ' calibrated vol ', pow(sig, 2)
        print 'vol for generating data ', pow(self.sigmas, 2)


class CalSurface(CalVolTestSuite):
    """ Test Calibrating volatility surface as in the paper. """

    def runTest(self):
        """ 1.- instantiate engine.
            2.- set calibration domain.
            3.- set calibration data.
            4.- set initial guess.
            5.- calibrate.
         """

        calib_vol_surface = CalibrateDVS()

        # define the grid for the finite element method.

        (time_grid, grid_price) = (self.ntime_steps, self.m_nsegs)
        calib_vol_surface.set_grid_spacing(grid_price, time_grid)

        # calibration data

        (ki, si, r) = (self.m_strikes, self.m_prices, self.m_r)
        (ti, style) = (self.m_t2mat, self.m_cstyles)
        calib_vol_surface.set_cal_data(ki, si, r, ti, style, self.m_contract)

        # constant inital volatility.

        calib_vol_surface.set_initial_guess([.35])

        calb_vol_curve.set_cal_method('fmin')
        (time_2m, sig) = calib_vol_surface.calibrate()


class CalibVolSurfaceCCE(CalVolTestSuite):
    """ calibrate volatility surface solving a pde fix point problem. """

    @classmethod
    def vol_exact_scalar(cls, stock_price):
        """ exact solution. """

        return 2. + 0.1 * exp( -0.2 * stock_price[0] )

    @classmethod
    def vol_exact(cls, stock_price):
        """ exact solution. """

        tmp = stock_price[:, 0].reshape((stock_price[:, 0].size, 1))

        # matplot tends to behave wier when too much precision is
        # pump in and the range of variation is small
        return 2. + 0.1 * exp( -0.2 * tmp )

    def setUp(self):
        """ initialize all the needed variables. """

        self.m_r = float(0.05)
        self.m_price_enigine = vanillaCallPrice
        self.m_vega = vegaVanilla
        self.m_stock_prices = asarray(linspace(60, 100, 10))
        self.m_strikes = asarray(linspace(50, 90, 10))
        self.m_t2m = 10 * ones(self.m_strikes.size) / 254.  # 10 days

        self.m_t2m[3:6] *= 2.

        self.m_t2m[7:10] *= 3.

        # generate call prices:

        self.m_prices = empty(self.m_strikes.shape)
        for (i, val) in izip(count(), c_[c_[self.m_stock_prices,
                             self.m_strikes], self.m_t2m]):
            (price, strike, t2m) = val
            self.m_prices[i] = self.m_price_enigine([price], strike, self.m_r,
                    self.vol_exact_scalar(array([price, t2m])), t2m)

    def runTest(self):

        # (x, y) ~ (stoke price, time change)
        # set boundary conditions

        (smin, smax) = (min(self.m_stock_prices) - 10.0,
                        max(self.m_stock_prices) + 10.0)
        (tmin, tmax) = (min(self.m_t2m) - 0.020, max(self.m_t2m)
                        + 0.020)
        (delta_x, delta_y) = (1.5, 0.009)

        calisigma = CalibVolSurfCCE()
        calisigma.set_simu_domain(smin, tmin, smax, tmax)

        calisigma.set_bott_top(self.vol_exact, self.vol_exact)
        calisigma.set_xlxr(self.vol_exact, self.vol_exact)

        calisigma.set_interest_rate(self.m_r)
        calisigma.set_dx_dy(delta_x, delta_y)

        # TODO REMOVE load will be reset at each iteration
        # calisigma.set_load(None)

        # set market data, and analytic option pricer.
        calisigma.set_calibration_data(self.m_strikes, self.m_t2m,
                self.m_prices)
        calisigma.set_analytical_enginee(self.m_price_enigine,
                self.m_vega)

        # Calibration Loop parameters.

        (tol, niter) = (1.0e-06, 5)
        calisigma.set_tol_and_niter(tol, niter)
        calisigma.set_relax_param(.1)

        # inital Guess.
        calisigma.set_initial_guess(0.2)

        # Calibrate the model

        (x_grid, y_grid, sigma) = calisigma.calibrate()

        # Post process results.

        (x_mesh, y_mesh) = meshgrid(x_grid, y_grid)
        exactsigma = self.vol_exact(c_[x_mesh.flatten(), y_mesh.flatten()])
        error_l2 = norm(sigma.flatten() - exactsigma)

        plot_results = True
        if plot_results:
            plot_engine = DerivativePlots()
            plot_engine.set_vol_surface(x_mesh, y_mesh, sigma,
                                 exactsigma.reshape(sigma.shape))

            plot_engine.plot_vol_surfaces()
        print ' Error : ', error_l2
        assert error_l2 < 1.0e-06


def my_suite():

    suite = unittest.TestSuite()

    suite.addTest(CalCurve())
    #suite.addTest(CalSurface())
    #suite.addTest(CalibVolSurfaceCCE())

    return suite


if __name__ == '__main__':

    runGlobalLogger()
    runner = unittest.TextTestRunner(verbosity=2)
    result = runner.run(my_suite())

